import mxnet as mx
from mxnet import nd, gluon
import time

data_shape = [(3, 3, 16, 224, 224), (3, 3, 128, 128, 128)]

for shape in data_shape:
	data = mx.random.uniform(shape=shape)
	weight_shape = (32, shape[1], 3, 3, 3)
	weight = mx.nd.ones(shape=weight_shape)

	num_iter = 10
	dry_run = 5
	for i in range(num_iter):
		if i == dry_run:
			tic = time.time()
		out = mx.nd.Convolution(data, weight, kernel=(3,3,3), stride=(1,1,1), num_filter=weight_shape[0], pad=(0,0,0), dilate=(2,2,2), no_bias=True, cudnn_off=True, name='conv3d')
		out.asnumpy()
	print("For shape : {}".format(shape))
	print("Average time cost is %f sec" % ((time.time() - tic)/(num_iter-dry_run))

import mxnet as mx
import time

def perf_conv3d():
    shapes = [(3, 3, 12, 128, 64), (3, 3, 12, 128, 128), (3, 3, 128, 128, 128), (5, 5, 16, 128, 128), (5, 5, 128, 128, 128)]
    kernel_shapes = [(2, 2, 2), (3, 3, 3)]
    num_filter = 32

    warmup = 10
    runs = 40

    for shape in shapes:
        for kernel_shape in kernel_shapes:
            a = mx.random.uniform(shape=shape)
            w = mx.random.uniform(shape=(num_filter, shape[1], kernel_shape[0], kernel_shape[1], kernel_shape[2]))

            tic = 0
            for i in range(runs + warmup):
                if i == warmup:
                    tic = time.time()
                _ = mx.nd.Convolution(data=a, weight=w, kernel=kernel_shape, num_filter=num_filter, stride=(1,1,1), pad=(0,0,0), no_bias=True, cudnn_off=True)
                mx.nd.waitall()

            toc = time.time()
            print('conv3d benchmark, shape={}, kernel_shape={} time={} ms/iter'.format(shape, kernel_shape, (toc-tic)*1000/(runs-warmup)))


if __name__ == '__main__':
    perf_conv3d()